#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PBPK Modeling Test Script using 7.10 parameter data

This script demonstrates PBPK modeling using the comprehensive test data
from 7_10_pbpk_test_data.json with complete visualization.

Author: PBPK Modeling System
Version: 1.0.0
Date: 2025-07-23
"""

import json
import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, Any, List
import os
import warnings
warnings.filterwarnings('ignore')

class PBPKModelTest:
    """PBPK modeling test class using 7.10 parameter data"""
    
    def __init__(self, test_data_path: str = None):
        """Initialize with test data"""
        if test_data_path is None:
            # Use absolute path based on script location
            current_dir = os.path.dirname(os.path.abspath(__file__))
            test_data_path = os.path.join(current_dir, "..", "data", "sample", "7_10_pbpk_test_data.json")
        
        self.test_data_path = test_data_path
        self.data = self._load_test_data()
        self.time_points = None
        self.results = {}
    
    def _load_test_data(self) -> Dict[str, Any]:
        """Load test data from JSON file"""
        try:
            with open(self.test_data_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except FileNotFoundError:
            print(f"Test data file not found: {self.test_data_path}")
            raise
    
    def calculate_pk_parameters(self, drug_name: str, subject_type: str = "adult_male") -> Dict[str, Any]:
        """Calculate PK parameters for a specific drug"""
        
        # Get drug parameters
        physico = self.data["physicochemical"][drug_name]
        absorption = self.data["absorption"][drug_name]
        demographics = self.data["demographics"][subject_type]
        
        # Calculate PK parameters
        dose = 200  # mg (standard test dose)
        
        # Volume of distribution based on tissue composition
        v_plasma = demographics["weight"] * 0.045  # L
        v_liver = self.data["organ_volumes"]["liver"]["volume"] * demographics["weight"] / 70
        v_kidney = self.data["organ_volumes"]["kidney"]["volume"] * demographics["weight"] / 70
        v_muscle = self.data["organ_volumes"]["muscle"]["volume"] * demographics["weight"] / 70
        v_adipose = self.data["organ_volumes"]["adipose"]["volume"] * demographics["weight"] / 70
        
        # Calculate tissue:plasma partition coefficients
        k_pt = {}
        for tissue, comp in self.data["tissue_composition"].items():
            # Use Poulin & Theil method
            f_water = comp["water_fraction"]
            f_lipid = comp["lipid_fraction"]
            f_protein = comp["protein_fraction"]
            
            k_pt[tissue] = (
                f_water + 
                f_lipid * 10**physico["logp"] + 
                f_protein * 0.5
            )
        
        # Calculate clearance
        total_clearance = absorption["clearance"] * (demographics["weight"] / 70)
        renal_clearance = absorption["renal_clearance"] * (demographics["weight"] / 70)
        hepatic_clearance = absorption["hepatic_clearance"] * (demographics["weight"] / 70)
        
        # Calculate bioavailability
        f_abs = absorption["bioavailability"]
        
        return {
            "drug_name": drug_name,
            "subject_type": subject_type,
            "dose": dose,
            "clearance_total": total_clearance,
            "clearance_renal": renal_clearance,
            "clearance_hepatic": hepatic_clearance,
            "volume_central": v_plasma,
            "volume_peripheral": v_muscle + v_adipose,
            "volume_liver": v_liver,
            "volume_kidney": v_kidney,
            "bioavailability": f_abs,
            "k_elimination": total_clearance / (v_plasma + v_liver),
            "partition_coefficients": k_pt,
            "half_life": 0.693 * (v_plasma + v_liver) / total_clearance
        }
    
    def run_pbpk_simulation(self, drug_name: str, time_span: float = 24.0, 
                          subject_type: str = "adult_male") -> Dict[str, Any]:
        """Run complete PBPK simulation"""
        
        # Get PK parameters
        pk_params = self.calculate_pk_parameters(drug_name, subject_type)
        
        # Setup time points
        self.time_points = np.linspace(0, time_span, 100)
        dt = self.time_points[1] - self.time_points[0]
        
        # Initialize compartments
        dose = pk_params["dose"]
        f_abs = pk_params["bioavailability"]
        
        # Central compartment (plasma + liver)
        v_central = pk_params["volume_central"] + pk_params["volume_liver"]
        
        # Peripheral compartments
        v_muscle = self.data["organ_volumes"]["muscle"]["volume"] * 15.0  # Adjusted for 70kg
        v_adipose = self.data["organ_volumes"]["adipose"]["volume"] * 15.0
        
        # Rate constants
        k_elim = pk_params["k_elimination"]
        k_dist_muscle = 0.1  # Distribution rate
        k_dist_adipose = 0.05
        
        # Simulation arrays
        plasma_conc = np.zeros(len(self.time_points))
        liver_conc = np.zeros(len(self.time_points))
        muscle_conc = np.zeros(len(self.time_points))
        adipose_conc = np.zeros(len(self.time_points))
        
        # Initial conditions
        plasma_conc[0] = dose * f_abs / v_central
        
        # Absorption phase (first-order)
        ka = self.data["absorption"][drug_name]["ka"]
        absorbed = np.zeros(len(self.time_points))
        
        for i in range(1, len(self.time_points)):
            # Absorption
            if i == 1:
                absorbed[i] = dose * f_abs * (1 - np.exp(-ka * self.time_points[i]))
            
            # Central compartment
            d_plasma = (-k_elim * plasma_conc[i-1] - 
                       k_dist_muscle * plasma_conc[i-1] - 
                       k_dist_adipose * plasma_conc[i-1])
            
            plasma_conc[i] = max(0, plasma_conc[i-1] + d_plasma * dt)
            
            # Peripheral compartments
            muscle_conc[i] = muscle_conc[i-1] + (k_dist_muscle * plasma_conc[i-1] - 
                                               0.1 * muscle_conc[i-1]) * dt
            adipose_conc[i] = adipose_conc[i-1] + (k_dist_adipose * plasma_conc[i-1] - 
                                                  0.05 * adipose_conc[i-1]) * dt
            liver_conc[i] = plasma_conc[i] * pk_params["partition_coefficients"]["liver"]
        
        return {
            "drug_name": drug_name,
            "subject_type": subject_type,
            "time_points": self.time_points,
            "plasma_concentration": plasma_conc,
            "liver_concentration": liver_conc,
            "muscle_concentration": muscle_conc,
            "adipose_concentration": adipose_conc,
            "pk_parameters": pk_params,
            "cmax": np.max(plasma_conc),
            "tmax": self.time_points[np.argmax(plasma_conc)],
            "auc": np.trapezoid(plasma_conc, self.time_points)
        }
    
    def run_multi_drug_analysis(self, drug_list: List[str] = None) -> Dict[str, Any]:
        """Run analysis for multiple drugs"""
        
        if drug_list is None:
            drug_list = list(self.data["physicochemical"].keys())
        
        results = {}
        for drug in drug_list:
            results[drug] = self.run_pbpk_simulation(drug)
        
        return results
    
    def generate_plots(self, results: Dict[str, Any], save_path: str = "../output"):
        """Generate comprehensive visualization"""
        
        os.makedirs(save_path, exist_ok=True)
        
        # Create figure with subplots
        fig = plt.figure(figsize=(20, 15))
        
        # 1. Plasma concentration-time curves
        ax1 = plt.subplot(3, 3, 1)
        colors = plt.cm.Set3(np.linspace(0, 1, len(results)))
        
        for i, (drug, result) in enumerate(results.items()):
            ax1.plot(result["time_points"], result["plasma_concentration"], 
                    label=drug, color=colors[i], linewidth=2)
        
        ax1.set_xlabel('Time (hours)')
        ax1.set_ylabel('Plasma Concentration (mg/L)')
        ax1.set_title('Plasma Concentration-Time Curves')
        ax1.legend()
        ax1.grid(True, alpha=0.3)
        
        # 2. Organ concentration comparison
        ax2 = plt.subplot(3, 3, 2)
        for i, (drug, result) in enumerate(results.items()):
            final_conc = {
                'Plasma': result["plasma_concentration"][-1],
                'Liver': result["liver_concentration"][-1],
                'Muscle': result["muscle_concentration"][-1],
                'Adipose': result["adipose_concentration"][-1]
            }
            ax2.bar([f"{k}\n{drug}" for k in final_conc.keys()], 
                   final_conc.values(), alpha=0.7, label=drug)
        
        ax2.set_ylabel('Final Concentration (mg/L)')
        ax2.set_title('Organ Concentration Comparison')
        ax2.legend()
        plt.xticks(rotation=45)
        
        # 3. PK parameters comparison
        ax3 = plt.subplot(3, 3, 3)
        drugs = list(results.keys())
        cmax_values = [results[d]["cmax"] for d in drugs]
        auc_values = [results[d]["auc"] for d in drugs]
        
        x = np.arange(len(drugs))
        width = 0.35
        
        ax3.bar(x - width/2, cmax_values, width, label='Cmax', alpha=0.8)
        ax3.bar(x + width/2, [auc/10 for auc in auc_values], width, label='AUC/10', alpha=0.8)
        
        ax3.set_xlabel('Drugs')
        ax3.set_ylabel('Values')
        ax3.set_title('PK Parameters Comparison')
        ax3.set_xticks(x)
        ax3.set_xticklabels(drugs, rotation=45)
        ax3.legend()
        
        # 4. Tissue distribution over time
        ax4 = plt.subplot(3, 3, 4)
        sample_drug = list(results.keys())[0]
        ax4.plot(results[sample_drug]["time_points"], 
                results[sample_drug]["plasma_concentration"], 
                label='Plasma', linewidth=2)
        ax4.plot(results[sample_drug]["time_points"], 
                results[sample_drug]["liver_concentration"], 
                label='Liver', linewidth=2)
        ax4.plot(results[sample_drug]["time_points"], 
                results[sample_drug]["muscle_concentration"], 
                label='Muscle', linewidth=2)
        ax4.plot(results[sample_drug]["time_points"], 
                results[sample_drug]["adipose_concentration"], 
                label='Adipose', linewidth=2)
        
        ax4.set_xlabel('Time (hours)')
        ax4.set_ylabel('Concentration (mg/L)')
        ax4.set_title(f'Tissue Distribution - {sample_drug}')
        ax4.legend()
        ax4.grid(True, alpha=0.3)
        
        # 5. BCS classification scatter plot
        ax5 = plt.subplot(3, 3, 5)
        x_vals = []
        y_vals = []
        labels = []
        colors_bcs = {'I': 'green', 'II': 'orange', 'III': 'blue', 'IV': 'red'}
        
        for drug in results.keys():
            drug_data = self.data["physicochemical"][drug]
            x_vals.append(drug_data["solubility"])
            y_vals.append(drug_data["permeability"] * 1e6)
            labels.append(drug)
            
            ax5.scatter(drug_data["solubility"], drug_data["permeability"] * 1e6,
                       c=colors_bcs[drug_data["bcs_class"]], s=100, alpha=0.7)
            ax5.annotate(drug, (drug_data["solubility"], drug_data["permeability"] * 1e6))
        
        ax5.set_xlabel('Solubility (mg/mL)')
        ax5.set_ylabel('Permeability (×10⁻⁶ cm/s)')
        ax5.set_title('BCS Classification')
        ax5.set_xscale('log')
        ax5.set_yscale('log')
        ax5.grid(True, alpha=0.3)
        
        # 6. Half-life comparison
        ax6 = plt.subplot(3, 3, 6)
        half_lives = [results[d]["pk_parameters"]["half_life"] for d in results.keys()]
        bars = ax6.bar(results.keys(), half_lives, color=colors[:len(results)], alpha=0.7)
        ax6.set_ylabel('Half-life (hours)')
        ax6.set_title('Elimination Half-life Comparison')
        plt.xticks(rotation=45)
        
        # Add value labels on bars
        for bar, hl in zip(bars, half_lives):
            height = bar.get_height()
            ax6.text(bar.get_x() + bar.get_width()/2., height + 0.1,
                    f'{hl:.1f}h', ha='center', va='bottom')
        
        # 7. Clearance pathways
        ax7 = plt.subplot(3, 3, 7)
        clearance_data = []
        for drug in results.keys():
            pk = results[drug]["pk_parameters"]
            clearance_data.append({
                'drug': drug,
                'renal': pk["clearance_renal"],
                'hepatic': pk["clearance_hepatic"]
            })
        
        bottom = np.zeros(len(results))
        x_pos = np.arange(len(results))
        
        ax7.bar(x_pos, [d['renal'] for d in clearance_data], 
               label='Renal', alpha=0.7, color='skyblue')
        ax7.bar(x_pos, [d['hepatic'] for d in clearance_data], 
               bottom=[d['renal'] for d in clearance_data],
               label='Hepatic', alpha=0.7, color='lightcoral')
        
        ax7.set_xlabel('Drugs')
        ax7.set_ylabel('Clearance (L/h)')
        ax7.set_title('Clearance Pathways')
        ax7.set_xticks(x_pos)
        ax7.set_xticklabels(results.keys(), rotation=45)
        ax7.legend()
        
        # 8. Volume of distribution
        ax8 = plt.subplot(3, 3, 8)
        volumes = []
        for drug in results.keys():
            pk = results[drug]["pk_parameters"]
            volumes.append({
                'drug': drug,
                'central': pk["volume_central"],
                'peripheral': pk["volume_peripheral"]
            })
        
        ax8.bar(x_pos, [v['central'] for v in volumes], label='Central', alpha=0.7)
        ax8.bar(x_pos, [v['peripheral'] for v in volumes], 
               bottom=[v['central'] for v in volumes],
               label='Peripheral', alpha=0.7)
        
        ax8.set_xlabel('Drugs')
        ax8.set_ylabel('Volume (L)')
        ax8.set_title('Volume of Distribution')
        ax8.set_xticks(x_pos)
        ax8.set_xticklabels(results.keys(), rotation=45)
        ax8.legend()
        
        # 9. Summary statistics
        ax9 = plt.subplot(3, 3, 9)
        ax9.axis('off')
        
        summary_text = "PBPK Modeling Summary\n\n"
        for drug, result in results.items():
            summary_text += f"{drug}:\n"
            summary_text += f"  Cmax: {result['cmax']:.2f} mg/L\n"
            summary_text += f"  Tmax: {result['tmax']:.1f} h\n"
            summary_text += f"  AUC: {result['auc']:.1f} mg·h/L\n"
            summary_text += f"  t1/2: {result['pk_parameters']['half_life']:.1f} h\n\n"
        
        ax9.text(0.1, 0.9, summary_text, transform=ax9.transAxes, 
                fontsize=10, verticalalignment='top', 
                bbox=dict(boxstyle="round,pad=0.3", facecolor="lightgray"))
        
        plt.tight_layout()
        plt.savefig(f"{save_path}/pbpk_modeling_comprehensive.png", dpi=300, bbox_inches='tight')
        plt.show()
        
        return f"{save_path}/pbpk_modeling_comprehensive.png"

def main():
    """Main test function"""
    print("=" * 60)
    print("PBPK Modeling Test Script")
    print("Using 7.10 Parameter Data")
    print("=" * 60)
    
    # Initialize model
    model = PBPKModelTest()
    
    # Run analysis for all drugs
    print("\n1. Running PBPK simulations...")
    results = model.run_multi_drug_analysis()
    
    print("\n2. Results Summary:")
    print("-" * 40)
    for drug, result in results.items():
        print(f"{drug.upper()}:")
        print(f"  Cmax: {result['cmax']:.2f} mg/L")
        print(f"  Tmax: {result['tmax']:.1f} hours")
        print(f"  AUC: {result['auc']:.1f} mg·h/L")
        print(f"  Half-life: {result['pk_parameters']['half_life']:.1f} hours")
        print()
    
    print("3. Generating comprehensive visualizations...")
    chart_path = model.generate_plots(results)
    print(f"📊 Charts saved to: {chart_path}")
    
    print("\n4. Detailed PK Parameters:")
    print("-" * 40)
    for drug, result in results.items():
        pk = result['pk_parameters']
        print(f"\n{drug.upper()} PK Parameters:")
        print(f"  Total Clearance: {pk['clearance_total']:.2f} L/h")
        print(f"  Renal Clearance: {pk['clearance_renal']:.2f} L/h")
        print(f"  Hepatic Clearance: {pk['clearance_hepatic']:.2f} L/h")
        print(f"  Volume Central: {pk['volume_central']:.1f} L")
        print(f"  Volume Peripheral: {pk['volume_peripheral']:.1f} L")
    
    print("\n" + "=" * 60)
    print("✅ PBPK Modeling Test Complete!")
    print("=" * 60)

if __name__ == "__main__":
    main()